import pandas as pd
import argparse
from pathlib import Path

# 默认搜索字符串数组
DEFAULT_SEARCH_STRINGS = [
"2c-car-service"
]

def filter_data(input_file, output_file, search_strings, column_name):
    # 根据文件扩展名选择读取方式
    file_extension = Path(input_file).suffix.lower()
    
    try:
        if file_extension == '.csv':
            df = pd.read_csv(input_file)
        elif file_extension in ('.xlsx', '.xls'):
            df = pd.read_excel(input_file)
        else:
            raise ValueError("不支持的文件格式，请提供 .csv 或 .xlsx/.xls 文件")
        
        # 检查列是否存在
        if column_name not in df.columns:
            available_columns = ", ".join(df.columns)
            raise KeyError(f"列 '{column_name}' 不存在。可用列: {available_columns}")
        
        # 过滤指定列中包含任意一个搜索字符串的行
        if not search_strings:
            filtered_df = df.iloc[0:0]  # 如果搜索字符串为空，返回空DataFrame
        else:
            filtered_df = df[df[column_name].str.contains('|'.join(search_strings), case=False, na=False)]
        
        # 输出结果到新文件
        output_extension = Path(output_file).suffix.lower()
        if output_extension == '.csv':
            filtered_df.to_csv(output_file, index=False)
        else:
            filtered_df.to_excel(output_file, index=False)
        
        print(f"过滤完成！输入: {len(df)} 行 → 输出: {len(filtered_df)} 行")
        print(f"结果已保存到: {output_file}")
    
    except Exception as e:
        print(f"错误: {e}")
        print("\n提示:")
        print("1. 检查输入文件是否存在且格式正确")
        print(f"2. 确认文件包含列 '{column_name}'")
        print(f"3. 当前文件可用列: {', '.join(df.columns) if 'df' in locals() else '未知'}")

# 执行命令 python src\excel.py -i src\服务接口.csv -o output-2025081552.csv
if __name__ == "__main__":
    # 设置命令行参数
    parser = argparse.ArgumentParser(description="过滤文件指定列中包含关键字的行（支持.csv/.xlsx）")
    parser.add_argument("-i", "--input", required=True, help="输入文件路径（.csv 或 .xlsx）")
    parser.add_argument("-o", "--output", required=True, help="输出文件路径（.csv 或 .xlsx）")
    parser.add_argument(
        "-s", "--search",
        nargs="*",
        default=DEFAULT_SEARCH_STRINGS,
        help=f"要匹配的字符串列表（空格分隔，默认值：{DEFAULT_SEARCH_STRINGS}）"
    )
    parser.add_argument(
        "-c", "--column",
        default="应用",
        help="要过滤的列名（默认：'B'）"
    )
    
    args = parser.parse_args()
    filter_data(args.input, args.output, args.search, args.column)